from PIL import Image
import numpy as np
import os
import matplotlib
matplotlib.use('Agg')  # 设置为非交互式后端
from tqdm import tqdm

# 1. 加载原始测试图像（确保与裁剪时的尺寸一致）
def load_test_image(img_path, target_size=(224, 224)):
    """加载并预处理测试图像"""
    if not os.path.exists(img_path):
        raise FileNotFoundError(f"图像路径不存在: {img_path}")
    
    image = Image.open(img_path).convert("RGB")
    # 调整为模型输入尺寸（224×224）
    image = image.resize(target_size)
    return image

# 2. 裁剪所有子图（使用修复后的无resize逻辑）
def crop_subimage(image, total_h, total_w, x, y):
    """基础子图裁剪函数（16×16）"""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    width, height = image.size
    sub_width = width // total_w  # 224//14=16
    sub_height = height // total_h  # 224//14=16
    
    left = x * sub_width
    upper = y * sub_height
    right = left + sub_width
    lower = upper + sub_height
    
    return image.crop((left, upper, right, lower))

def crop_all_subimages_by_ids(image, total_h, total_w, ids=None):
    """裁剪指定ID的子图，返回{id: 子图}的字典（便于后续按顺序拼接）"""
    subimg_dict = []  # key: curr_ids, value: 子图
    curr_ids = 0
    
    # 行优先遍历（与模型patch顺序一致：y从0→13，x从0→13）
    for y in range(total_h):
        for x in range(total_w):
            if ids is None or curr_ids in ids:    
                subimage = crop_subimage(image, total_h, total_w, x, y)
                subimg_dict.append(subimage)
            curr_ids += 1
    return subimg_dict

# 修改后的子图裁剪函数
def crop_all_siglip_subimages_by_ids(image, total_h, total_w, ids=None, target_size=384, patch_size=14):
    """裁剪指定ID的子图，返回子图列表
    Args:
        image: 原始图像
        total_h: 高度方向的patch数量
        total_w: 宽度方向的patch数量
        ids: 要裁剪的子图ID列表，None表示裁剪所有
        target_size: 模型期望的输入尺寸 (384)
        patch_size: 模型使用的patch尺寸 (14)
    """
    subimg_dict = []
    curr_ids = 0
    
    # 获取原始图像尺寸
    orig_width, orig_height = image.size
    
    # 计算resize 384x384后的实际感受野尺寸
    # 对于384x384的图像，每个patch的物理尺寸是14x14
    # 在原始图像上，等效的感受野尺寸需要按比例缩放
    scale_w = orig_width / target_size
    scale_h = orig_height / target_size
    
    # 每个子图在原始图像上的实际尺寸
    sub_width = int(patch_size * scale_w)
    sub_height = int(patch_size * scale_h)
    
    # 行优先遍历（与模型patch顺序一致）
    for y in range(total_h):
        for x in range(total_w):
            if ids is None or curr_ids in ids:
                # 计算原始图像上的裁剪坐标
                # 注意：这里需要考虑图像边缘的对齐问题
                left = int(x * patch_size * scale_w)
                upper = int(y * patch_size * scale_h)
                right = min(left + sub_width, orig_width)
                lower = min(upper + sub_height, orig_height)
                
                # 裁剪子图
                subimage = image.crop((left, upper, right, lower))
                subimage = subimage.resize((target_size, target_size))
                subimg_dict.append(subimage)
            curr_ids += 1
    
    return subimg_dict

# 3. 将子图按原始顺序拼接回完整图像
def stitch_subimages(subimg_dict, grid_size=(14, 14), subimg_size=(16, 16), target_size=(224, 224)):
    """
    按行优先顺序拼接子图
    Args:
        subimg_dict: {curr_ids: 子图} 字典（curr_ids对应模型patch索引）
        grid_size: 子图网格（14行14列，与模型patch数量一致）
        subimg_size: 单张子图尺寸（16×16）
        target_size: 拼接后的目标尺寸（224×224）
    Returns:
        拼接后的完整图像
    """
    # 创建空白画布（RGB模式）
    stitched_image = Image.new('RGB', target_size)
    
    total_rows, total_cols = grid_size
    sub_w, sub_h = subimg_size
    
    # 按行优先顺序粘贴子图（与裁剪顺序完全一致）
    for curr_ids in tqdm(range(total_rows * total_cols), desc="拼接子图"):
        if curr_ids not in subimg_dict:
            raise KeyError(f"缺失子图ID: {curr_ids}，请检查裁剪逻辑")
        
        subimg = subimg_dict[curr_ids]
        
        # 计算当前子图在画布上的坐标（行优先）
        row = curr_ids // total_cols  # 行索引（0~13）
        col = curr_ids % total_cols   # 列索引（0~13）
        
        # 计算粘贴位置（左上角坐标）
        left = col * sub_w
        upper = row * sub_h
        
        # 将子图粘贴到画布
        stitched_image.paste(subimg, (left, upper))
    
    return stitched_image

# 4. 完整验证流程
if __name__ == "__main__":
    # 配置参数（与模型和裁剪逻辑匹配）
    IMG_PATH = "data/image.png"  # 你的测试图像路径
    GRID_SIZE = (14, 14)        # 14×14子图（对应模型196个patch）
    SUBIMG_SIZE = (16, 16)      # 单张子图尺寸（与模型patch_size一致）
    TARGET_SIZE = (224, 224)    # 目标拼接尺寸（与模型输入一致）
    ALL_IDS = list(range(14*14)) # 所有子图ID（0~195）
    
    # 步骤1：加载原始图像
    original_image = load_test_image(IMG_PATH, TARGET_SIZE)
    original_image.save("visualizations/original_image.png")
    print("原始图像已保存至: visualizations/original_image.png")
    
    # 步骤2：裁剪所有子图
    subimg_dict = crop_all_subimages_by_ids(original_image, GRID_SIZE[0], GRID_SIZE[1], ALL_IDS)
    print(f"成功裁剪 {len(subimg_dict)} 个子图（应等于196）")
    
    # 步骤3：拼接子图
    stitched_image = stitch_subimages(subimg_dict, GRID_SIZE, SUBIMG_SIZE, TARGET_SIZE)
    
    # 步骤4：保存并对比结果
    os.makedirs("visualizations", exist_ok=True)
    stitched_image.save("visualizations/stitched_image.png")
    print("拼接后的图像已保存至: visualizations/stitched_image.png")
    
    # 验证拼接是否完全一致（可选，通过像素差异判断）
    original_np = np.array(original_image)
    stitched_np = np.array(stitched_image)
    pixel_diff = np.sum(np.abs(original_np - stitched_np))
    
    if pixel_diff == 0:
        print("✅ 拼接验证通过！子图顺序与原始图像完全一致，无区域错位")
    else:
        print(f"⚠️  拼接验证警告：原始图像与拼接图像存在 {pixel_diff} 个像素差异（可能是裁剪/粘贴误差）")